Skip to content

feat: kernel-density CDF meshes RectangularKernelAdaptDensity/Image - #374

Merged
Jammy2211 merged 3 commits into
mainfrom
feature/rectangular-kernel-cdf-mesh
Jul 10, 2026
Merged

Jammy2211 merged 3 commits into
mainfrom
feature/rectangular-kernel-cdf-mesh

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Adds opt-in kernel-density-CDF variants of the adaptive rectangular meshes: RectangularKernelAdaptDensity / RectangularKernelAdaptImage (issue #373). The per-axis transform is the smooth kernel CDF F(x) = Σᵢ wᵢ·Φ((x−xᵢ)/h) (Enzi et al. arXiv:2606.30620 RTU formulation) instead of the empirical point-rank CDF, which the 2026-07-09 JAX gradient audit proved makes the likelihood exactly piecewise-constant in mass/shear at pixelization over-sampling 1 and on the interferometer sparse path.

Properties by construction: strictly monotone (no jitter hack), C^∞ in interp queries and traced-point positions (no ranks, no sorts, no argsort anywhere), duplicate-safe (no 1/Δknot terms). The forward transform is evaluated exactly and maps the data bounding box onto the unit square (matching the linear convention); the inverse is a fixed n_knots table lookup whose queries are unit-grid constants, so gradients flow smoothly through the table values. bandwidth defaults to the mesh pixel scale and is exposed as a mesh kwarg. The linear and spline meshes are untouched.

Function-level JAX sanity (local, x64): eager == JIT to 1e-13; AD grad matches central FD to ~5e-10 relative through both the bilinear-weights path and the inverse-table path.

API Changes

Added only — two new opt-in mesh classes (aa.mesh.RectangularKernelAdaptDensity, aa.mesh.RectangularKernelAdaptImage), their interpolator, and additive default-None kwargs on AbstractMeshGeometry. No existing symbol removed, renamed, or changed in behaviour.
See full details below.

Test Plan

  • pytest test_autoarray/ — 892 passed (16 new tests in test_rectangular_kernel.py: monotonicity, bounding-box mapping, round-trip, duplicate safety, large-bandwidth limit, mapper shape/consistency vs linear, areas totals vs linear, interpolator/geometry plumbing)
  • Function-level JAX check: eager-vs-JIT consistency + AD-vs-FD on forward and inverse paths
  • Certification (workspace leg, gating merge): autolens_workspace_test/scripts/jax_grad/imaging_pixelization.py kernel variants at os_pix=1 AND 4 with strict FD on ALL params; interferometer.py sparse path; eager-vs-JIT; figure-of-merit parity within a few e-4 of the linear meshes
Full API Changes (for automation & release notes)

Added

  • aa.mesh.RectangularKernelAdaptDensity(shape=, bandwidth=1.0, n_knots=64) — density-adaptive rectangular mesh with kernel-density CDF transform (opt-in).
  • aa.mesh.RectangularKernelAdaptImage(shape=, weight_power=1.0, weight_floor=0.0, bandwidth=1.0, n_knots=64) — image-adaptive variant; adapt-image weights become the kernel weights.
  • autoarray.inversion.mesh.interpolator.rectangular_kernel — new module: create_transforms_kernel(), adaptive_rectangular_mappings_weights_via_interpolation_from_kernel(), adaptive_rectangular_areas_from_kernel(), adaptive_rectangular_transformed_grid_from_kernel(), InterpolatorRectangularKernel, KERNEL_CDF_DEFAULT_BANDWIDTH, KERNEL_CDF_DEFAULT_KNOTS.

Changed Signature (additive, backwards-compatible)

  • AbstractMeshGeometry.__init__(..., kernel_bandwidth=None, kernel_knots=None) — new optional kwargs; MeshGeometryRectangular.areas_transformed / edges_transformed route through the kernel helpers when kernel_bandwidth is set (mirrors the existing spline_deg dispatch).

Migration

  • None required — all changes are additive and opt-in.

Coordination note

Adjacent to #372 (rect-adapt, parked at sign-off): both touch mesh_geometry/rectangular.py in distinct hunks (its edges_transformed body vs this PR's dispatch branches) — clean merge in either order; if #372 lands first the kernel path inherits the corrected edges automatically.

Merge gate: do not merge until the workspace-leg certification PR (jax_grad strict-FD assertions) is green — that is this task's success criterion.

🤖 Generated with Claude Code

…ants

Opt-in adaptive rectangular meshes whose per-axis CDF transform is a smooth
kernel-density CDF F(x) = sum_i w_i * Phi((x - x_i)/h) (Enzi et al.
arXiv:2606.30620 RTU formulation) instead of the empirical point-rank CDF.
Strictly monotone by construction, C-infinity in interp queries and traced
point positions (no ranks, no sorts), duplicate-safe — so mass/shear
gradients are smooth in every configuration, including imaging at
pixelization over-sampling 1 and the interferometer sparse path where the
empirical CDF makes the likelihood exactly piecewise-constant.

Forward transform is evaluated exactly (unit square <-> data bounding box,
matching the linear convention); inverse is a fixed n_knots table lookup
whose queries are unit-grid constants, so gradients flow through the table
values. Bandwidth defaults to the mesh pixel scale and is exposed as a mesh
kwarg. Linear and spline meshes are untouched.

Closes nothing; certification via autolens_workspace_test jax_grad scripts
follows on the workspace leg (#373).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 10, 2026
Jammy2211 and others added 2 commits July 10, 2026 10:25
The bilinear interpolation mapper places reconstruction value (row r,
col c) at node U_y=(n_y-r-1)/(n_y-3), U_x=(c-1)/(n_x-3) in unit CDF space
(guard ring at the border, rows flipped by the mapper's row = n - i), but
edges_transformed drew a uniform [0,1] partition — pcolormesh plots put
flux up to ~1.5 mesh pixels from where the mapper scatters it (the "small
offsets" of issue #372, confirmed by a delta-function reproduction). The
edges are now the node midpoints pushed through the same CDF machinery
(guard-safe denominator for degenerate n<=3 meshes; border guard cells
clamp to the data span). Both consumers inherit the fix: the pcolormesh
plot path and rectangular_rotated's warped-centre reconstruction; spline
and linear CDFs share the upstream index convention. Plot geometry only —
the mapper is untouched, so inversions and likelihoods are byte-identical.

#372

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Folded in feature/rect-adapt (c177391, PR #375) by merge — the kernel meshes' plot geometry inherits the #372 edge fix through the shared edges_transformed dispatch. Full suite re-run on the combined tree: 893 passed. Verification of the fix (incl. on the kernel meshes) is on #372. If #375 merges first this PR needs nothing; if this merges first, #375 becomes a no-op diff.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant